#Basic imports 
from numpy import *
from pylab import *
from FFT import *
from dislin import *
import wave
import sys
import struct

# open the wave file
fp = wave.open(sys.argv[1],"rb")
sample_rate = fp.getframerate()
total_num_samps = fp.getnframes()
fft_length = int(sys.argv[2])
num_fft = (total_num_samps / fft_length ) - 2

# create temporary working array
temp = zeros((num_fft,fft_length),Float)

# read in the data from the file
for i in range(num_fft):
    tempb = fp.readframes(fft_length);
    temp[i,:] = array(struct.unpack("%dB"%(fft_length), \
                tempb),Float) - 128.0
fp.close()

# Window the data
temp = temp * hamming(fft_length)

# Transform with the FFT, Return Power
freq_pwr  = 10*log10(1e-20+abs(real_fft(temp,fft_length)))

# Plot the result
n_out_pts = (fft_length / 2) + 1
y_axis = 0.5*float(sample_rate) / n_out_pts * \
         arange(n_out_pts)
x_axis = (total_num_samps / float(sample_rate)) / \
         num_fft * arange(num_fft)
setvar('X',"Time (sec)")
setvar('Y',"Frequency (Hertz)")
conshade(freq_pwr,x_axis,y_axis)
disfin()